# coding: utf-8

import xml.etree.ElementTree as ET
import os
import glob
import cv2

names_dict = {}
cnt = 0
f = open('../data/face_mask/face_mask.names', 'r').readlines()
for line in f:
    line = line.strip()
    names_dict[line] = cnt
    cnt += 1

def parse_xml(path):
    tree = ET.parse(path)
    img_name = path.split('/')[-1][:-4]

    height = tree.findtext("./size/height")
    width = tree.findtext("./size/width")

    if height==None or width==None or height=='0' or width == '0':
        img = cv2.imread(path.replace(".xml", ".jpg"))
        height, width, _ = img.shape

    objects = [img_name, str(width), str(height)]

    for obj in tree.findall('object'):
        difficult = obj.find('difficult').text
        if difficult == '1':
            continue
        name = obj.find('name').text
        bbox = obj.find('bndbox')
        xmin = bbox.find('xmin').text
        ymin = bbox.find('ymin').text
        xmax = bbox.find('xmax').text
        ymax = bbox.find('ymax').text

        name = str(names_dict[name])
        objects.extend([name, xmin, ymin, xmax, ymax])
    if len(objects) > 3:
        return objects
    else:
        return None


test_cnt = 0


def gen_test_txt(txt_path):
    global test_cnt
    f = open(txt_path, 'w')

    img_names = glob.glob(f"{val_path}/*.xml")
    for img_name in img_names:
        img_name = img_name.strip()
        xml_path = img_name
        objects = parse_xml(xml_path)
        if objects:
            objects[0] = img_name.replace(".xml", ".jpg")
            if os.path.exists(objects[0]):
                objects.insert(0, str(test_cnt))
                test_cnt += 1
                objects = ' '.join(objects) + '\n'
                f.write(objects)
    f.close()


train_cnt = 0


def gen_train_txt(txt_path):
    global train_cnt
    f = open(txt_path, 'w')

    img_names = glob.glob(f"{train_path}/*.xml")
    for img_name in img_names:
        img_name = img_name.strip()
        xml_path = img_name
        objects = parse_xml(xml_path)
        if objects:
            objects[0] = img_name.replace(".xml", ".jpg")
            if os.path.exists(objects[0]):
                objects.insert(0, str(train_cnt))
                train_cnt += 1
                objects = ' '.join(objects) + '\n'
                f.write(objects)
    f.close()

train_path = "/home/yangna/data/face_mask/train"        # 1.
val_path = "/home/yangna/data/face_mask/val"            # 2.

gen_train_txt('../data/face_mask/train.txt')            # 3
gen_test_txt('../data/face_mask/val.txt')               # 4

